# -*- coding: utf-8 -*-
"""
@Time    : 2024/8/5 14:00 
@Author  : ZhangShenao 
@File    : 2.MultiQuery多查询策略.py
@Desc    : 使用MultiQueryRetriever实现MultiQuery多查询策略

多查询策略也被称为子查询,是一种用于生成子问题的技术其核心思想是:
在问答过程中,为了更好地理解和回答主问题,系统会自动生成并提出与主问题相关的子问题
这些子问题通常具有更具体的细节,可以帮助大语言模型更深入地理解主问题,从而进行更加准确的检索并提供正确的答案

多查询策略是最基础、最简单的RAG优化优化,不涉及到复杂的逻辑与算法,会稍微影响单次对话的耗时
并且由于需要转换query一般较小,以及生成sub-queries时对LLM的能力要求并不高
因此在实际的LLM应用开发中,通常使用参数较小的本地模型+针对性优化的Prompt即可完成任务

而且为了减少模型的幻觉,一般都将temperature设置为0,确保生成的文本更加有确定性

MultiQueryRetriever采用的生成相似提问的Prompt默认是英文的
在中文场景下,可以替换成自定义的中文的Prompt

"""
import dotenv
from langchain.retrievers import MultiQueryRetriever
from langchain_community.chat_models import ChatZhipuAI
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_core.prompts import ChatPromptTemplate
from langchain_pinecone import PineconeVectorStore

dotenv.load_dotenv()

# 基于Pinecone,创建基础的VectorStoreRetriever
embeddings = ZhipuAIEmbeddings(model="embedding-3")
vector_store = PineconeVectorStore(
    index_name="rag",  # 指定索引名称
    embedding=embeddings,  # 指定Embedding模型
    namespace="llm-ops",  # 指定Namespace
    text_key="text",  # 指定文本属性key
)
base_retriever = vector_store.as_retriever(
    search_type="similarity_score_threshold",
    search_kwargs={
        "k": 3,
        "score_threshold": 0.5
    }
)

# 创建MultiQueryRetriever多查询检索器
prompt = ChatPromptTemplate.from_template(
    "你是一个AI语言模型助手。你的任务是生成给定用户问题的3个不同版本，以从向量数据库中检索相关文档。"
    "通过提供用户问题的多个视角，你的目标是帮助用户克服基于距离的相似性搜索的一些限制。"
    "请用换行符分隔这些替代问题。"
    "原始问题：{question}"
)
multi_query_retriever = MultiQueryRetriever.from_llm(
    retriever=base_retriever,  # 指定底层进行检索的Prompt
    llm=ChatZhipuAI(model="glm-4-flash", temperature=0),  # 指定用于生成相似提问的LLM,可以将temperature设置为0,降低模型出现幻觉的概率
    prompt=prompt,  # MultiQueryRetriever采用的生成相似提问的Prompt默认是英文的，可以替换成中文的Prompt
)

# 执行检索
# MultiQueryRetriever底层就是基于原始提问,生成3个相似的提问,然后分别进行检索,最后将结果Document进行合并去重
docs = multi_query_retriever.invoke("关于LLM应用配置的文档有哪些？")
print(len(docs))
print(docs[0])
